import pandas as pd
from sklearn.preprocessing import LabelEncoder

def preprocess_data(file_name, include_renewal):
    # 读取Excel文件
    df = pd.read_excel(file_name)

    # 变量相关性分析
    # 对分类变量进行编码处理
    cat_cols = df.select_dtypes(include=['object']).columns

    # policy_term字段：去掉'年'字，并转换为整数
    df['policy_term'] = df['policy_term'].str.replace('年', '').astype(int)

    # 对income_level进行数值编码
    income_map = {'低': 0, '中': 1, '高': 2}
    df['income_level'] = df['income_level'].map(lambda x: income_map.get(x))
    # 对education_level进行数值编码
    education_map = {'高中': 0, '本科': 1, '硕士': 2, '博士': 3}

    df['education_level'] = df['education_level'].map(lambda x: education_map.get(x))
    # 对marital_status进行数值编码
    marital_map = {'单身': 0, '已婚': 1, '离异': 2}
    df['marital_status'] = df['marital_status'].map(lambda x: marital_map.get(x))

    # 对gender进行数值编码
    gender_map = {'女': 0, '男': 1}
    df['gender'] = df['gender'].map(lambda x: gender_map.get(x))


    # 对claim_history进行数值编码
    claim_history_map = {'是': 1, '否': 0}
    df['claim_history'] = df['claim_history'].map(lambda x: claim_history_map.get(x))

    # 若不包含renewal字段，不进行编码处理
    if not include_renewal:
        df = df.drop(columns=['renewal'], errors='ignore')
    else:
        # 对renewal进行数值编码
        renewal_map = {'Yes': 1, 'No': 0}
        df['renewal'] = df['renewal'].map(lambda x: renewal_map.get(x))

    # 将分类变量转换为数值型
    le = LabelEncoder()
    cate_cols = ['birth_region', 'insurance_region', 
    'occupation', 'policy_type']
    for col in cate_cols:
        df[col] = le.fit_transform(df[col])

    # 对日期类型字段进行处理，将其转换为时间戳
    for col in df.select_dtypes(include=['datetime']).columns:
        df[col] = pd.to_datetime(df[col])
        df[col] = df[col].astype('int64') 

    # 转换为自某一起始点的天数
    base_date = pd.to_datetime('2010-01-01')
    # 确保日期列是日期时间类型
    if pd.api.types.is_numeric_dtype(df['policy_start_date']):
        df['policy_start_date'] = pd.to_datetime(df['policy_start_date'])
    if pd.api.types.is_numeric_dtype(df['policy_end_date']):
        df['policy_end_date'] = pd.to_datetime(df['policy_end_date'])
    df['policy_start_date_day'] = (df['policy_start_date'] - base_date).dt.days
    df['policy_end_date_day'] = (df['policy_end_date'] -
    base_date).dt.days

    return df

if __name__ == '__main__':
    file_name = 'data/policy_data.xlsx'
    include_renewal = True
    processed_df = preprocess_data(file_name, include_renewal)
    processed_df.to_excel('data/processed_data_train.xlsx', index=False)
    
    file_name = 'data/policy_test.xlsx'
    include_renewal = False
    processed_df = preprocess_data(file_name, include_renewal)
    processed_df.to_excel('data/processed_data_test.xlsx', index=False)